#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
小红书数据采集系统数据库管理模块
提供SQLite数据库的增删改查操作
"""

import sqlite3
import datetime
import json
from typing import Dict, List, Union, Optional, Any
import os
import threading
import re

class XHSDataDB:
    """小红书数据采集系统数据库管理类"""
    
    def __init__(self, db_path='xhs_database.db'):
        """
        初始化数据库连接
        
        Args:
            db_path: SQLite数据库文件路径
        """
        # 确保数据库文件路径正确
        # 如果是相对路径，则从当前工作目录解析
        if not os.path.isabs(db_path):
            # 获取当前文件所在目录
            current_dir = os.path.dirname(os.path.abspath(__file__))
            self.db_path = os.path.join(current_dir, db_path)
        else:
            # 如果是绝对路径，直接使用
            self.db_path = db_path
            
        print(f"使用数据库路径: {self.db_path}")
        
        self.conn = None
        self.cursor = None
        self._thread_id = threading.get_ident()
        self._connect()
    
    def _connect(self):
        """建立数据库连接"""
        # 确保在同一线程中进行连接操作
        current_thread = threading.get_ident()
        if self.conn is not None and self._thread_id != current_thread:
            # 如果是不同线程，关闭旧连接并重新连接
            self.close()
            
        # 创建新连接
        self.conn = sqlite3.connect(self.db_path)
        # 设置行工厂，返回字典形式的结果
        self.conn.row_factory = sqlite3.Row
        self.cursor = self.conn.cursor()
        self._thread_id = current_thread
    
    def close(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
            self.cursor = None
        if self.conn:
            self.conn.close()
            self.conn = None
    
    # 确保每个方法开始时都先检查连接有效性
    def _ensure_connection(self):
        """确保连接有效，如果无效则重新连接"""
        current_thread = threading.get_ident()
        if self.conn is None or self._thread_id != current_thread:
            self._connect()

    ###########################################
    # 采集任务相关操作
    ###########################################
    
    def get_all_tasks(self, limit=100, offset=0) -> List[Dict]:
        """
        获取所有采集任务
        
        Args:
            limit: 限制返回数量
            offset: 跳过记录数
            
        Returns:
            采集任务记录列表
        """
        self._ensure_connection()
        self.cursor.execute(
            "SELECT * FROM collection_tasks ORDER BY task_id DESC LIMIT ? OFFSET ?", 
            (limit, offset)
        )
        return [dict(row) for row in self.cursor.fetchall()]
    
    def get_task_by_id(self, task_id) -> Optional[Dict]:
        """
        根据ID获取采集任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            采集任务记录或None
        """
        self._ensure_connection()
        self.cursor.execute("SELECT * FROM collection_tasks WHERE task_id = ?", (task_id,))
        row = self.cursor.fetchone()
        return dict(row) if row else None
    
    def add_task(self, task_data: Dict) -> int:
        """
        添加采集任务
        
        Args:
            task_data: 任务数据字典
            
        Returns:
            新添加任务的ID
        """
        self._ensure_connection()
        columns = ', '.join(task_data.keys())
        placeholders = ', '.join(['?' for _ in task_data])
        values = tuple(task_data.values())
        
        self.cursor.execute(
            f"INSERT INTO collection_tasks ({columns}) VALUES ({placeholders})",
            values
        )
        self.conn.commit()
        return self.cursor.lastrowid
    
    def update_task(self, task_id, update_data: Dict) -> bool:
        """
        更新采集任务
        
        Args:
            task_id: 任务ID
            update_data: 更新的数据
            
        Returns:
            更新是否成功
        """
        self._ensure_connection()
        set_clause = ', '.join([f"{column} = ?" for column in update_data.keys()])
        values = list(update_data.values())
        values.append(task_id)
        
        self.cursor.execute(
            f"UPDATE collection_tasks SET {set_clause} WHERE task_id = ?",
            tuple(values)
        )
        self.conn.commit()
        return self.cursor.rowcount > 0
    
    def delete_task(self, task_id) -> bool:
        """
        删除采集任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            删除是否成功
        """
        self._ensure_connection()
        self.cursor.execute("DELETE FROM collection_tasks WHERE task_id = ?", (task_id,))
        self.conn.commit()
        return self.cursor.rowcount > 0
    
    ###########################################
    # 笔记(文章)相关操作
    ###########################################
    
    def get_all_notes(self, limit=100, offset=0, keyword=None) -> List[Dict]:
        """
        获取所有笔记
        
        Args:
            limit: 限制返回数量
            offset: 跳过记录数
            keyword: 关键词筛选
            
        Returns:
            笔记记录列表
        """
        self._ensure_connection()
        query = "SELECT * FROM notes"
        params = []
        
        if keyword:
            query += " WHERE keyword LIKE ?"
            params.append(f"%{keyword}%")
        
        query += " ORDER BY id DESC LIMIT ? OFFSET ?"
        params.extend([limit, offset])
        
        self.cursor.execute(query, tuple(params))
        return [dict(row) for row in self.cursor.fetchall()]
    
    def get_note_by_id(self, note_id) -> Optional[Dict]:
        """
        根据ID获取笔记
        
        Args:
            note_id: 笔记ID
            
        Returns:
            笔记记录或None
        """
        self._ensure_connection()
        self.cursor.execute("SELECT * FROM notes WHERE id = ?", (note_id,))
        row = self.cursor.fetchone()
        return dict(row) if row else None
    
    def get_notes_by_keyword(self, keyword, limit=100) -> List[Dict]:
        """
        根据关键词获取笔记
        
        Args:
            keyword: 关键词
            limit: 限制返回数量
            
        Returns:
            笔记记录列表
        """
        self._ensure_connection()
        self.cursor.execute(
            "SELECT * FROM notes WHERE keyword LIKE ? ORDER BY id DESC LIMIT ?", 
            (f"%{keyword}%", limit)
        )
        return [dict(row) for row in self.cursor.fetchall()]
    
    def add_note(self, note_data: Dict) -> int:
        """
        添加笔记
        
        Args:
            note_data: 笔记数据
            
        Returns:
            新添加笔记的ID
        """
        # 如果没有设置采集时间，设置为当前时间
        if 'collect_time' not in note_data:
            note_data['collect_time'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        self._ensure_connection()
        columns = ', '.join(note_data.keys())
        placeholders = ', '.join(['?' for _ in note_data])
        values = tuple(note_data.values())
        
        try:
            self.cursor.execute(
                f"INSERT INTO notes ({columns}) VALUES ({placeholders})",
                values
            )
            self.conn.commit()
            return self.cursor.lastrowid
        except sqlite3.IntegrityError:
            # 处理唯一约束冲突（例如URL重复）
            if 'url' in note_data:
                self.cursor.execute("SELECT id FROM notes WHERE url = ?", (note_data['url'],))
                row = self.cursor.fetchone()
                if row:
                    return row['id']
            return -1
    
    def update_note(self, note_id, update_data: Dict) -> bool:
        """
        更新笔记信息
        
        Args:
            note_id: 笔记ID
            update_data: 更新的数据
            
        Returns:
            更新是否成功
        """
        self._ensure_connection()
        set_clause = ', '.join([f"{column} = ?" for column in update_data.keys()])
        values = list(update_data.values())
        values.append(note_id)
        
        self.cursor.execute(
            f"UPDATE notes SET {set_clause} WHERE id = ?",
            tuple(values)
        )
        self.conn.commit()
        return self.cursor.rowcount > 0
    
    def delete_note(self, note_id) -> bool:
        """
        删除笔记
        
        Args:
            note_id: 笔记ID
            
        Returns:
            删除是否成功
        """
        self._ensure_connection()
        
        # 获取笔记URL (用于可能的关联删除)
        self.cursor.execute("SELECT url FROM notes WHERE id = ?", (note_id,))
        note = self.cursor.fetchone()
        
        if not note:
            return False
            
        # 开始事务
        self.conn.execute("BEGIN TRANSACTION")
        
        try:
            # 删除相关评论
            self.cursor.execute("DELETE FROM comments WHERE article_url = ?", (note['url'],))
            
            # 删除笔记
            self.cursor.execute("DELETE FROM notes WHERE id = ?", (note_id,))
            
            self.conn.commit()
            return True
        except Exception as e:
            self.conn.rollback()
            print(f"删除笔记失败: {e}")
            return False
    
    ###########################################
    # 评论相关操作
    ###########################################
    
    def get_comments_by_note(self, note_id, limit=100, offset=0) -> List[Dict]:
        """
        获取笔记的所有评论
        
        Args:
            note_id: 笔记ID
            limit: 限制返回数量
            offset: 跳过记录数
            
        Returns:
            评论记录列表
        """
        self._ensure_connection()
        
        # 先获取笔记的URL
        self.cursor.execute("SELECT url FROM notes WHERE id = ?", (note_id,))
        note = self.cursor.fetchone()
        
        if not note:
            return []
            
        # 根据笔记URL查询评论
        self.cursor.execute(
            "SELECT * FROM comments WHERE article_url = ? ORDER BY id DESC LIMIT ? OFFSET ?", 
            (note['url'], limit, offset)
        )
        return [dict(row) for row in self.cursor.fetchall()]
    
    def add_comment(self, comment_data: Dict) -> int:
        """
        添加评论
        
        Args:
            comment_data: 评论数据
            
        Returns:
            新添加评论的ID
        """
        # 如果没有设置采集时间，设置为当前时间
        if 'collect_time' not in comment_data:
            comment_data['collect_time'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        self._ensure_connection()
        columns = ', '.join(comment_data.keys())
        placeholders = ', '.join(['?' for _ in comment_data])
        values = tuple(comment_data.values())
        
        self.cursor.execute(
            f"INSERT INTO comments ({columns}) VALUES ({placeholders})",
            values
        )
        
        # 更新笔记的评论计数
        if 'note_id' in comment_data:
            self.cursor.execute(
                "UPDATE notes SET comments_count = comments_count + 1 WHERE id = ?", 
                (comment_data['note_id'],)
            )
        
        self.conn.commit()
        return self.cursor.lastrowid
    
    ###########################################
    # 用户相关操作
    ###########################################
    
    def get_all_users(self, limit=100, offset=0) -> List[Dict]:
        """
        获取所有用户
        
        Args:
            limit: 限制返回数量
            offset: 跳过记录数
            
        Returns:
            用户记录列表
        """
        self._ensure_connection()
        self.cursor.execute(
            "SELECT * FROM users ORDER BY followers_count DESC LIMIT ? OFFSET ?", 
            (limit, offset)
        )
        return [dict(row) for row in self.cursor.fetchall()]
    
    def get_user_by_id(self, user_id) -> Optional[Dict]:
        """
        根据ID获取用户
        
        Args:
            user_id: 用户ID
            
        Returns:
            用户记录或None
        """
        self._ensure_connection()
        self.cursor.execute("SELECT * FROM users WHERE user_id = ?", (user_id,))
        row = self.cursor.fetchone()
        return dict(row) if row else None
    
    def add_or_update_user(self, user_data: Dict) -> str:
        """
        添加或更新用户
        
        Args:
            user_data: 用户数据
            
        Returns:
            用户ID
        """
        # 检查用户是否已存在
        user_id = user_data.get('user_id')
        if not user_id:
            return ""
        
        existing_user = self.get_user_by_id(user_id)
        
        if existing_user:
            # 更新用户
            update_data = {k: v for k, v in user_data.items() if k != 'user_id'}
            if 'last_updated' not in update_data:
                update_data['last_updated'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
            self._ensure_connection()
            set_clause = ', '.join([f"{column} = ?" for column in update_data.keys()])
            values = list(update_data.values())
            values.append(user_id)
            
            self.cursor.execute(
                f"UPDATE users SET {set_clause} WHERE user_id = ?",
                tuple(values)
            )
        else:
            # 添加用户
            if 'first_seen_date' not in user_data:
                user_data['first_seen_date'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            if 'last_updated' not in user_data:
                user_data['last_updated'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
            self._ensure_connection()
            columns = ', '.join(user_data.keys())
            placeholders = ', '.join(['?' for _ in user_data])
            values = tuple(user_data.values())
            
            self.cursor.execute(
                f"INSERT INTO users ({columns}) VALUES ({placeholders})",
                values
            )
        
        self.conn.commit()
        return user_id
    
    ###########################################
    # 位置相关操作
    ###########################################
    
    def get_all_locations(self, limit=100, offset=0) -> List[Dict]:
        """
        获取所有位置
        
        Args:
            limit: 限制返回数量
            offset: 跳过记录数
            
        Returns:
            位置记录列表
        """
        self._ensure_connection()
        self.cursor.execute(
            "SELECT * FROM locations ORDER BY count DESC LIMIT ? OFFSET ?", 
            (limit, offset)
        )
        return [dict(row) for row in self.cursor.fetchall()]
    
    def add_or_update_location(self, location_data: Dict) -> int:
        """
        添加或更新位置
        
        Args:
            location_data: 位置数据
            
        Returns:
            位置ID
        """
        location_name = location_data.get('location_name')
        if not location_name:
            return -1
        
        # 检查位置是否已存在
        self._ensure_connection()
        self.cursor.execute("SELECT location_id FROM locations WHERE location_name = ?", (location_name,))
        row = self.cursor.fetchone()
        
        if row:
            # 更新位置
            location_id = row['location_id']
            update_data = {k: v for k, v in location_data.items() if k != 'location_name'}
            if 'last_updated' not in update_data:
                update_data['last_updated'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            if 'count' not in update_data:
                update_data['count'] = 1
            else:
                # 如果提供了count但要增加而不是设置
                self.cursor.execute("UPDATE locations SET count = count + ? WHERE location_id = ?", 
                                    (update_data['count'], location_id))
                del update_data['count']
            
            if update_data:  # 如果还有其他字段需要更新
                set_clause = ', '.join([f"{column} = ?" for column in update_data.keys()])
                values = list(update_data.values())
                values.append(location_id)
                
                self.cursor.execute(
                    f"UPDATE locations SET {set_clause} WHERE location_id = ?",
                    tuple(values)
                )
            
            self.conn.commit()
            return location_id
        else:
            # 添加位置
            if 'last_updated' not in location_data:
                location_data['last_updated'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            if 'count' not in location_data:
                location_data['count'] = 1
            
            self._ensure_connection()
            columns = ', '.join(location_data.keys())
            placeholders = ', '.join(['?' for _ in location_data])
            values = tuple(location_data.values())
            
            self.cursor.execute(
                f"INSERT INTO locations ({columns}) VALUES ({placeholders})",
                values
            )
            self.conn.commit()
            return self.cursor.lastrowid
    
    ###########################################
    # 统计分析
    ###########################################
    
    def get_keyword_statistics(self) -> List[Dict]:
        """
        获取关键词统计
        
        Returns:
            关键词统计列表 [{'keyword': '关键词', 'count': 数量}, ...]
        """
        self._ensure_connection()
        self.cursor.execute(
            "SELECT keyword, COUNT(*) as count FROM notes WHERE keyword IS NOT NULL AND keyword != '' GROUP BY keyword ORDER BY count DESC"
        )
        return [dict(row) for row in self.cursor.fetchall()]
    
    def get_note_types_statistics(self, limit=10) -> List[Dict]:
        """
        获取笔记类型(keyword)及其对应的笔记数量
        
        Args:
            limit: 限制返回数量，默认10条
            
        Returns:
            类型统计列表 [{'keyword': '类型', 'count': 数量}, ...]
        """
        self._ensure_connection()
        self.cursor.execute(
            "SELECT keyword, COUNT(*) as count FROM notes WHERE keyword IS NOT NULL AND keyword != '' GROUP BY keyword ORDER BY count DESC LIMIT ?",
            (limit,)
        )
        return [dict(row) for row in self.cursor.fetchall()]
    
    def get_locations_from_edit_time(self, limit=20, offset=0, search=None) -> List[Dict]:
        """
        从notes表的edit_time字段中提取位置信息并进行统计
        
        Args:
            limit: 限制返回数量，默认20条
            offset: 偏移量，用于分页
            search: 搜索关键词，用于筛选位置
            
        Returns:
            位置统计列表 [{'location_name': '位置名称', 'count': 数量, 'location_id': 序号}, ...]
        """
        self._ensure_connection()
        
        # 使用正则表达式提取位置信息
        # SQLite不支持强大的正则功能，所以我们获取所有edit_time然后在Python中处理
        query = "SELECT edit_time FROM notes WHERE edit_time IS NOT NULL AND edit_time != ''"
        params = []
        
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        
        # 统计位置信息
        location_counts = {}
        
        for row in rows:
            edit_time = row['edit_time']
            # 正则表达式匹配日期后的位置信息：匹配格式如 "04-21 四川" 中的 "四川"
            match = re.search(r'\d{2}-\d{2}(?:\s+编辑于)?\s+([^\d\s].*?)$', edit_time)
            if match:
                location = match.group(1).strip()
                if location:
                    if search and search.lower() not in location.lower():
                        continue
                    
                    if location in location_counts:
                        location_counts[location] += 1
                    else:
                        location_counts[location] = 1
        
        # 按数量排序并限制返回结果
        sorted_locations = sorted(
            [{'location_name': loc, 'count': count} for loc, count in location_counts.items()],
            key=lambda x: x['count'], 
            reverse=True
        )
        
        # 应用分页并添加序号作为location_id
        paginated_locations = sorted_locations[offset:offset+limit]
        for i, loc in enumerate(paginated_locations, start=offset+1):
            loc['location_id'] = i
        
        # 返回总数以便分页
        total_count = len(sorted_locations)
        
        return paginated_locations, total_count
    
    def get_location_statistics(self) -> List[Dict]:
        """
        获取位置统计
        
        Returns:
            位置统计列表
        """
        # 使用从edit_time提取位置的方法,限制最多20条
        locations, _ = self.get_locations_from_edit_time(limit=20)
        return locations
    
    def get_dashboard_statistics(self) -> Dict:
        """
        获取仪表盘统计数据
        
        Returns:
            统计数据字典
        """
        self._ensure_connection()
        stats = {}
        
        # 笔记总数
        self.cursor.execute("SELECT COUNT(*) as count FROM notes")
        stats['total_notes'] = self.cursor.fetchone()['count']
        
        # 评论总数
        self.cursor.execute("SELECT COUNT(*) as count FROM comments")
        stats['total_comments'] = self.cursor.fetchone()['count']
        
        # 用户总数 - 从评论表的user_id字段统计不重复用户数量
        try:
            self.cursor.execute("SELECT COUNT(DISTINCT user_id) as count FROM comments WHERE user_id IS NOT NULL AND user_id != ''")
            stats['total_users'] = self.cursor.fetchone()['count']
        except sqlite3.OperationalError:
            # 表可能不存在或其他错误
            stats['total_users'] = 0
        
        # 任务总数
        try:
            self.cursor.execute("SELECT COUNT(*) as count FROM collection_tasks")
            stats['total_tasks'] = self.cursor.fetchone()['count']
        except sqlite3.OperationalError:
            # 表可能不存在
            stats['total_tasks'] = 0
        
        # 最近任务
        try:
            self.cursor.execute("SELECT * FROM collection_tasks ORDER BY task_id DESC LIMIT 5")
            stats['recent_tasks'] = [dict(row) for row in self.cursor.fetchall()]
        except sqlite3.OperationalError:
            # 表可能不存在
            stats['recent_tasks'] = []
        
        # 热门笔记(点赞最多)
        self.cursor.execute("SELECT * FROM notes ORDER BY likes DESC LIMIT 5")
        stats['popular_notes'] = [dict(row) for row in self.cursor.fetchall()]
        
        return stats
    
    ###########################################
    # 导入导出功能
    ###########################################
    
    def export_to_json(self, file_path: str) -> bool:
        """
        导出数据到JSON文件
        
        Args:
            file_path: 文件路径
            
        Returns:
            导出是否成功
        """
        try:
            export_data = {
                'collection_tasks': [],
                'notes': [],
                'comments': [],
                'users': [],
                'locations': []
            }
            
            # 获取所有采集任务
            self._ensure_connection()
            self.cursor.execute("SELECT * FROM collection_tasks")
            export_data['collection_tasks'] = [dict(row) for row in self.cursor.fetchall()]
            
            # 获取所有笔记
            self._ensure_connection()
            self.cursor.execute("SELECT * FROM notes")
            export_data['notes'] = [dict(row) for row in self.cursor.fetchall()]
            
            # 获取所有评论
            self._ensure_connection()
            self.cursor.execute("SELECT * FROM comments")
            export_data['comments'] = [dict(row) for row in self.cursor.fetchall()]
            
            # 获取所有用户
            self._ensure_connection()
            self.cursor.execute("SELECT * FROM users")
            export_data['users'] = [dict(row) for row in self.cursor.fetchall()]
            
            # 获取所有位置
            self._ensure_connection()
            self.cursor.execute("SELECT * FROM locations")
            export_data['locations'] = [dict(row) for row in self.cursor.fetchall()]
            
            # 写入JSON文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(export_data, f, ensure_ascii=False, indent=2)
            
            return True
        except Exception as e:
            print(f"导出失败: {e}")
            return False
    
    def import_from_json(self, file_path: str) -> bool:
        """
        从JSON文件导入数据
        
        Args:
            file_path: 文件路径
            
        Returns:
            导入是否成功
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                import_data = json.load(f)
            
            # 开始事务
            self.conn.execute("BEGIN TRANSACTION")
            
            # 导入采集任务
            if 'collection_tasks' in import_data:
                for task in import_data['collection_tasks']:
                    self.add_task(task)
            
            # 导入笔记
            if 'notes' in import_data:
                for note in import_data['notes']:
                    self.add_note(note)
            
            # 导入评论
            if 'comments' in import_data:
                for comment in import_data['comments']:
                    self.add_comment(comment)
            
            # 导入用户
            if 'users' in import_data:
                for user in import_data['users']:
                    self.add_or_update_user(user)
            
            # 导入位置
            if 'locations' in import_data:
                for location in import_data['locations']:
                    self.add_or_update_location(location)
            
            # 提交事务
            self.conn.commit()
            
            return True
        except Exception as e:
            # 回滚事务
            self.conn.rollback()
            print(f"导入失败: {e}")
            return False 